import numpy as np

from shapely.geometry import Polygon

from collections import deque

from .. import util
from .. import bounds
from .. import graph

from ..constants import tol_path as tol
from ..constants import log
from ..transformations import transform_points

from .simplify import fit_circle_check
from .traversal import resample_path

try:
    import networkx as nx
except BaseException as E:
    # create a dummy module which will raise the ImportError
    # or other exception only when someone tries to use networkx
    from ..exceptions import ExceptionModule
    nx = ExceptionModule(E)
try:
    from rtree import Rtree
except BaseException as E:
    # create a dummy module which will raise the ImportError
    from ..exceptions import ExceptionModule
    Rtree = ExceptionModule(E)


def enclosure_tree(polygons):
    """
    Given a list of shapely polygons with only exteriors,
    find which curves represent the exterior shell or root curve
    and which represent holes which penetrate the exterior.

    This is done with an R-tree for rough overlap detection,
    and then exact polygon queries for a final result.

    Parameters
    -----------
    polygons : (n,) shapely.geometry.Polygon
       Polygons which only have exteriors and may overlap

    Returns
    -----------
    roots : (m,) int
        Index of polygons which are root
    contains : networkx.DiGraph
       Edges indicate a polygon is
       contained by another polygon
    """
    tree = Rtree()
    # nodes are indexes in polygons
    contains = nx.DiGraph()
    for i, polygon in enumerate(polygons):
        # if a polygon is None it means creation
        # failed due to weird geometry so ignore it
        if polygon is None or len(polygon.bounds) != 4:
            continue
        # insert polygon bounds into rtree
        tree.insert(i, polygon.bounds)
        # make sure every valid polygon has a node
        contains.add_node(i)

    # loop through every polygon
    for i in contains.nodes():
        polygon = polygons[i]
        # we first query for bounding box intersections from the R-tree
        for j in tree.intersection(polygon.bounds):
            # if we are checking a polygon against itself continue
            if (i == j):
                continue
            # do a more accurate polygon in polygon test
            # for the enclosure tree information
            if polygons[i].contains(polygons[j]):
                contains.add_edge(i, j)
            elif polygons[j].contains(polygons[i]):
                contains.add_edge(j, i)

    # a root or exterior curve has an even number of parents
    # wrap in dict call to avoid networkx view
    degree = dict(contains.in_degree())

    # convert keys and values to numpy arrays
    indexes = np.array(list(degree.keys()))
    degrees = np.array(list(degree.values()))

    # roots are curves with an even inward degree (parent count)
    roots = indexes[(degrees % 2) == 0]

    # if there are multiple nested polygons split the graph
    # so the contains logic returns the individual polygons
    if len(degrees) > 0 and degrees.max() > 1:
        # collect new edges for graph
        edges = []
        # find edges of subgraph for each root and children
        for root in roots:
            children = indexes[degrees == degree[root] + 1]
            edges.extend(contains.subgraph(np.append(children, root)).edges())
        # stack edges into new directed graph
        contains = nx.from_edgelist(edges, nx.DiGraph())
        # if roots have no children add them anyway
        contains.add_nodes_from(roots)

    return roots, contains


def edges_to_polygons(edges, vertices):
    """
    Given an edge list of indices and associated vertices
    representing lines, generate a list of polygons.

    Parameters
    -----------
    edges : (n, 2) int
      Indexes of vertices which represent lines
    vertices : (m, 2) float
      Vertices in 2D space

    Returns
    ----------
    polygons : (p,) shapely.geometry.Polygon
      Polygon objects with interiors
    """

    # create closed polygon objects
    polygons = []
    # loop through a sequence of ordered traversals
    for dfs in graph.traversals(edges, mode='dfs'):
        try:
            # try to recover polygons before they are more complicated
            polygons.append(repair_invalid(Polygon(vertices[dfs])))
        except ValueError:
            continue

    # if there is only one polygon, just return it
    if len(polygons) == 1:
        return polygons

    # find which polygons contain which other polygons
    roots, tree = enclosure_tree(polygons)

    # generate list of polygons with proper interiors
    complete = []
    for root in roots:
        interior = list(tree[root].keys())
        shell = polygons[root].exterior.coords
        holes = [polygons[i].exterior.coords for i in interior]
        complete.append(Polygon(shell=shell,
                                holes=holes))
    return complete


def polygons_obb(polygons):
    """
    Find the OBBs for a list of shapely.geometry.Polygons
    """
    rectangles = [None] * len(polygons)
    transforms = [None] * len(polygons)
    for i, p in enumerate(polygons):
        transforms[i], rectangles[i] = polygon_obb(p)
    return np.array(transforms), np.array(rectangles)


def polygon_obb(polygon):
    """
    Find the oriented bounding box of a Shapely polygon.

    The OBB is always aligned with an edge of the convex hull of the polygon.

    Parameters
    -------------
    polygons : shapely.geometry.Polygon
      Input geometry

    Returns
    -------------
    transform : (3, 3) float
      Transformation matrix
      which will move input polygon from its original position
      to the first quadrant where the AABB is the OBB
    extents : (2,) float
      Extents of transformed polygon
    """
    if hasattr(polygon, 'exterior'):
        points = np.asanyarray(polygon.exterior.coords)
    elif isinstance(polygon, np.ndarray):
        points = polygon
    else:
        raise ValueError('polygon or points must be provided')
    return bounds.oriented_bounds_2D(points)


def transform_polygon(polygon, matrix):
    """
    Transform a polygon by a a 2D homogeneous transform.

    Parameters
    -------------
    polygon : shapely.geometry.Polygon
      2D polygon to be transformed.
    matrix  : (3, 3) float
      2D homogeneous transformation.

    Returns
    --------------
    result : shapely.geometry.Polygon
      Polygon transformed by matrix.
    """
    matrix = np.asanyarray(matrix, dtype=np.float64)

    if util.is_sequence(polygon):
        result = [transform_polygon(p, t)
                  for p, t in zip(polygon, matrix)]
        return result
    # transform the outer shell
    shell = transform_points(np.array(polygon.exterior.coords),
                             matrix)[:, :2]
    # transform the interiors
    holes = [transform_points(np.array(i.coords),
                              matrix)[:, :2]
             for i in polygon.interiors]
    # create a new polygon with the result
    result = Polygon(shell=shell, holes=holes)
    return result


def plot_polygon(polygon, show=True, **kwargs):
    """
    Plot a shapely polygon using matplotlib.

    Parameters
    ------------
    polygon : shapely.geometry.Polygon
      Polygon to be plotted
    show : bool
      If True will display immediately
    **kwargs
      Passed to plt.plot
    """
    import matplotlib.pyplot as plt

    def plot_single(single):
        plt.plot(*single.exterior.xy, **kwargs)
        for interior in single.interiors:
            plt.plot(*interior.xy, **kwargs)
    # make aspect ratio non- stupid
    plt.axes().set_aspect('equal', 'datalim')
    if util.is_sequence(polygon):
        [plot_single(i) for i in polygon]
    else:
        plot_single(polygon)

    if show:
        plt.show()


def resample_boundaries(polygon, resolution, clip=None):
    """
    Return a version of a polygon with boundaries resampled
    to a specified resolution.

    Parameters
    -------------
    polygon : shapely.geometry.Polygon
      Source geometry
    resolution : float
      Desired distance between points on boundary
    clip : (2,) int
      Upper and lower bounds to clip
      number of samples to avoid exploding count

    Returns
    ------------
    kwargs : dict
     Keyword args for a Polygon constructor `Polygon(**kwargs)`
    """
    def resample_boundary(boundary):
        # add a polygon.exterior or polygon.interior to
        # the deque after resampling based on our resolution
        count = boundary.length / resolution
        count = int(np.clip(count, *clip))
        return resample_path(boundary.coords, count=count)
    if clip is None:
        clip = [8, 200]
    # create a sequence of [(n,2)] points
    kwargs = {'shell': resample_boundary(polygon.exterior),
              'holes': deque()}
    for interior in polygon.interiors:
        kwargs['holes'].append(resample_boundary(interior))
    kwargs['holes'] = np.array(kwargs['holes'])
    return kwargs


def stack_boundaries(boundaries):
    """
    Stack the boundaries of a polygon into a single
    (n, 2) list of vertices.

    Parameters
    ------------
    boundaries : dict
      With keys 'shell', 'holes'

    Returns
    ------------
    stacked : (n, 2) float
      Stacked vertices
    """
    if len(boundaries['holes']) == 0:
        return boundaries['shell']
    result = np.vstack((boundaries['shell'],
                        np.vstack(boundaries['holes'])))
    return result


def medial_axis(polygon,
                resolution=None,
                clip=None):
    """
    Given a shapely polygon, find the approximate medial axis
    using a voronoi diagram of evenly spaced points on the
    boundary of the polygon.

    Parameters
    ----------
    polygon : shapely.geometry.Polygon
      The source geometry
    resolution : float
      Distance between each sample on the polygon boundary
    clip : None, or (2,) int
      Clip sample count to min of clip[0] and max of clip[1]

    Returns
    ----------
    edges : (n, 2) int
      Vertex indices representing line segments
      on the polygon's medial axis
    vertices : (m, 2) float
      Vertex positions in space
    """
    # a circle will have a single point medial axis
    if len(polygon.interiors) == 0:
        # what is the approximate scale of the polygon
        scale = np.reshape(polygon.bounds, (2, 2)).ptp(axis=0).max()
        # a (center, radius, error) tuple
        fit = fit_circle_check(
            polygon.exterior.coords, scale=scale)
        # is this polygon in fact a circle
        if fit is not None:
            # return an edge that has the center as the midpoint
            epsilon = np.clip(
                fit['radius'] / 500, 1e-5, np.inf)
            vertices = np.array(
                [fit['center'] + [0, epsilon],
                 fit['center'] - [0, epsilon]],
                dtype=np.float64)
            # return a single edge to avoid consumers needing to special case
            edges = np.array([[0, 1]], dtype=np.int64)
            return edges, vertices

    from scipy.spatial import Voronoi
    from shapely import vectorized

    if resolution is None:
        resolution = np.reshape(
            polygon.bounds, (2, 2)).ptp(axis=0).max() / 100

    # get evenly spaced points on the polygons boundaries
    samples = resample_boundaries(polygon=polygon,
                                  resolution=resolution,
                                  clip=clip)
    # stack the boundary into a (m,2) float array
    samples = stack_boundaries(samples)
    # create the voronoi diagram on 2D points
    voronoi = Voronoi(samples)
    # which voronoi vertices are contained inside the polygon
    contains = vectorized.contains(polygon, *voronoi.vertices.T)
    # ridge vertices of -1 are outside, make sure they are False
    contains = np.append(contains, False)
    # make sure ridge vertices is numpy array
    ridge = np.asanyarray(voronoi.ridge_vertices, dtype=np.int64)
    # only take ridges where every vertex is contained
    edges = ridge[contains[ridge].all(axis=1)]

    # now we need to remove uncontained vertices
    contained = np.unique(edges)
    mask = np.zeros(len(voronoi.vertices), dtype=np.int64)
    mask[contained] = np.arange(len(contained))

    # mask voronoi vertices
    vertices = voronoi.vertices[contained]
    # re-index edges
    edges_final = mask[edges]

    if tol.strict:
        # make sure we didn't screw up indexes
        assert (vertices[edges_final] -
                voronoi.vertices[edges]).ptp() < 1e-5

    return edges_final, vertices


def polygon_hash(polygon):
    """
    Return a vector containing values representitive of
    a particular polygon.

    Parameters
    ---------
    polygon : shapely.geometry.Polygon
      Input geometry

    Returns
    ---------
    hashed: (6), float
      Representitive values representing input polygon
    """
    result = np.array(
        [len(polygon.interiors),
         polygon.convex_hull.area,
         polygon.convex_hull.length,
         polygon.area,
         polygon.length,
         polygon.exterior.length],
        dtype=np.float64)
    return result


def random_polygon(segments=8, radius=1.0):
    """
    Generate a random polygon with a maximum number of sides and approximate radius.

    Parameters
    ---------
    segments : int
      The maximum number of sides the random polygon will have
    radius : float
      The approximate radius of the polygon desired

    Returns
    ---------
    polygon : shapely.geometry.Polygon
      Geometry object with random exterior and no interiors.
    """
    angles = np.sort(np.cumsum(np.random.random(
        segments) * np.pi * 2) % (np.pi * 2))
    radii = np.random.random(segments) * radius
    points = np.column_stack(
        (np.cos(angles), np.sin(angles))) * radii.reshape((-1, 1))
    points = np.vstack((points, points[0]))
    polygon = Polygon(points).buffer(0.0)
    if util.is_sequence(polygon):
        return polygon[0]
    return polygon


def polygon_scale(polygon):
    """
    For a Polygon object return the diagonal length of the AABB.

    Parameters
    ------------
    polygon : shapely.geometry.Polygon
      Source geometry

    Returns
    ------------
    scale : float
      Length of AABB diagonal
    """
    extents = np.reshape(polygon.bounds, (2, 2)).ptp(axis=0)
    scale = (extents ** 2).sum() ** .5

    return scale


def paths_to_polygons(paths, scale=None):
    """
    Given a sequence of connected points turn them into
    valid shapely Polygon objects.

    Parameters
    -----------
    paths : (n,) sequence
      Of (m, 2) float closed paths
    scale: float
      Approximate scale of drawing for precision

    Returns
    -----------
    polys : (p,) list
      Filled with Polygon or None

    """
    polygons = [None] * len(paths)
    for i, path in enumerate(paths):
        if len(path) < 4:
            # since the first and last vertices are identical in
            # a closed loop a 4 vertex path is the minimum for
            # non-zero area
            continue
        try:
            polygons[i] = repair_invalid(Polygon(path), scale)
        except ValueError:
            # raised if a polygon is unrecoverable
            continue
        except BaseException:
            log.error('unrecoverable polygon', exc_info=True)
    polygons = np.array(polygons)
    return polygons


def sample(polygon, count, factor=1.5, max_iter=10):
    """
    Use rejection sampling to generate random points inside a
    polygon.

    Parameters
    -----------
    polygon : shapely.geometry.Polygon
      Polygon that will contain points
    count : int
      Number of points to return
    factor : float
      How many points to test per loop
    max_iter : int
      Maximum number of intersection checks is:
      > count * factor * max_iter

    Returns
    -----------
    hit : (n, 2) float
      Random points inside polygon
      where n <= count
    """
    # do batch point-in-polygon queries
    from shapely import vectorized

    # get size of bounding box
    bounds = np.reshape(polygon.bounds, (2, 2))
    extents = bounds.ptp(axis=0)

    hit = []
    hit_count = 0
    per_loop = int(count * factor)

    for i in range(max_iter):
        # generate points inside polygons AABB
        points = np.random.random((per_loop, 2))
        points = (points * extents) + bounds[0]

        # do the point in polygon test and append resulting hits
        mask = vectorized.contains(polygon, *points.T)
        hit.append(points[mask])

        # keep track of how many points we've collected
        hit_count += len(hit[-1])

        # if we have enough points exit the loop
        if hit_count > count:
            break

    # stack the hits into an (n,2) array and truncate
    hit = np.vstack(hit)[:count]

    return hit


def repair_invalid(polygon, scale=None, rtol=.5):
    """
    Given a shapely.geometry.Polygon, attempt to return a
    valid version of the polygon through buffering tricks.

    Parameters
    -----------
    polygon : shapely.geometry.Polygon
      Source geometry
    rtol : float
      How close does a perimeter have to be
    scale : float or None
      For numerical precision reference

    Returns
    ----------
    repaired : shapely.geometry.Polygon
      Repaired polygon

    Raises
    ----------
    ValueError
      If polygon can't be repaired
    """
    if hasattr(polygon, 'is_valid') and polygon.is_valid:
        return polygon

    # basic repair involves buffering the polygon outwards
    # this will fix a subset of problems.
    basic = polygon.buffer(tol.zero)
    # if it returned multiple polygons check the largest
    if util.is_sequence(basic):
        basic = basic[np.argmax([i.area for i in basic])]

    # check perimeter of result against original perimeter
    if basic.is_valid and np.isclose(basic.length,
                                     polygon.length,
                                     rtol=rtol):
        return basic

    if scale is None:
        distance = 0.002 * polygon_scale(polygon)
    else:
        distance = 0.002 * scale

    # if there are no interiors, we can work with just the exterior
    # ring, which is often more reliable
    if len(polygon.interiors) == 0:
        # try buffering the exterior of the polygon
        # the interior will be offset by -tol.buffer
        rings = polygon.exterior.buffer(distance).interiors
        if len(rings) == 1:
            # reconstruct a single polygon from the interior ring
            recon = Polygon(shell=rings[0]).buffer(distance)
            # check perimeter of result against original perimeter
            if recon.is_valid and np.isclose(recon.length,
                                             polygon.length,
                                             rtol=rtol):
                return recon

        # try de-deuplicating the outside ring
        points = np.array(polygon.exterior)
        # remove any segments shorter than tol.merge
        # this is a little risky as if it was discretized more
        # finely than 1-e8 it may remove detail
        unique = np.append(True, (np.diff(points, axis=0)**2).sum(
            axis=1)**.5 > 1e-8)
        # make a new polygon with result
        dedupe = Polygon(shell=points[unique])
        # check result
        if dedupe.is_valid and np.isclose(dedupe.length,
                                          polygon.length,
                                          rtol=rtol):
            return dedupe

    # buffer and unbuffer the whole polygon
    buffered = polygon.buffer(distance).buffer(-distance)
    # if it returned multiple polygons check the largest
    if util.is_sequence(buffered):
        buffered = buffered[np.argmax([i.area for i in buffered])]
    # check perimeter of result against original perimeter
    if buffered.is_valid and np.isclose(buffered.length,
                                        polygon.length,
                                        rtol=rtol):
        log.debug('Recovered invalid polygon through double buffering')
        return buffered

    raise ValueError('unable to recover polygon!')
